import torch
import numpy as np


def time_adj_dp(data):
    dp = np.zeros((data.shape[-1], data.shape[-1]))
    # data的尺寸是(14, 4, seq)
    n = data.shape[0]
    res = np.zeros((n, n, data.shape[1]))
    for i in range(n):
        res[i, i] = 1
        for j in range(i + 1, n):
            for k in range(data.shape[1]):
                res[i, j, k] = time_adj_dp_func(data[i, k], data[j, k], dp)
            res[j, i] = res[i, j]
    return torch.from_numpy(res)


def time_adj_dp_func(t1, t2, dp):
    dp[0, 0] = (t1[0] - t2[0]) ** 2
    for i in range(1, len(t1)):
        dp[i, 0] = dp[i - 1, 0] + (t1[i] - t2[0]) ** 2
    for i in range(1, len(t2)):
        dp[0, i] = dp[0, i - 1] + (t1[0] - t2[i]) ** 2
    
    for i in range(1, len(t1)):
        for j in range(1, len(t2)):
            dp[i, j] = (t1[i] - t2[j]) ** 2
            dp[i, j] += min(dp[i - 1, j], dp[i, j - 1], dp[i - 1, j - 1])

    return dp[len(t1) - 1, len(t2) - 1]

if __name__ == "__main__":
    data = torch.randn(14, 4, 24)
    res = time_adj_dp(data)
    print(res.shape)